#include "ScanPrimitive.h"
#define GROUP_SIZE 128
#define GROUP_SIZE_STR "128"

ScanPrimitive::ScanPrimitive (UINT elemNum,
		QD3DResourcePack *p_dataPack, UINT structSize,
		const std::string &structBufType) : _elemNum (elemNum)
{
	QD3DUtils d3DUtils;

	_pGroupSizeDxConstBuf = d3DUtils.createConstBuffer< ConstVars > ();

	_tempPackPtrs.push_back (p_dataPack);
	for (; elemNum > GROUP_SIZE; ) {
		elemNum	= ceil (elemNum / (double)GROUP_SIZE);
		ID3D11Buffer *pTempDxBuf = d3DUtils.createStructuredBuffer (
				elemNum, structSize, NULL,
				D3D11_BIND_UNORDERED_ACCESS | D3D11_BIND_SHADER_RESOURCE);
		_tempPackPtrs.push_back (new QD3DResourcePack (pTempDxBuf));
	}

	QD3DShaderDesc scanDesc;
	scanDesc.setDefine ("TYPE", structBufType);
	scanDesc.setDefine ("GROUP_SIZE", GROUP_SIZE_STR);
	_pScanUpSweepCS = d3DUtils.createShaderFromFile< ID3D11ComputeShader > (
			QD3DUtils::COMPUTE_SHADER, "cs_5_0",
			"..\\Data\\Shaders\\Primitives\\ScanCS.hlsl", "scanUp", scanDesc);
	_pScanDownSweepCS =
			d3DUtils.createShaderFromFile< ID3D11ComputeShader > (
			QD3DUtils::COMPUTE_SHADER, "cs_5_0",
			"..\\Data\\Shaders\\Primitives\\ScanCS.hlsl", "scanDown", scanDesc);
}

ScanPrimitive::~ScanPrimitive ()
{
	SAFE_RELEASE (_pGroupSizeDxConstBuf);
	SAFE_RELEASE (_pScanUpSweepCS);
	SAFE_RELEASE (_pScanDownSweepCS);

	for (int i = 1; i < _tempPackPtrs.size (); ++i)
		delete _tempPackPtrs[i];
}

void ScanPrimitive::run (UINT elemNum)
{
	if (elemNum == -1)
		elemNum = _elemNum;

	QD3DUtils d3dUtils;
	ConstVars constVars;
	UINT *totalElemNumArr = new UINT[_tempPackPtrs.size ()];

	QVector< ID3D11Buffer * > dxConstBufPtrs (1);
	QVector< ID3D11ShaderResourceView * > dxSRVPtrs (1);
	QVector< ID3D11UnorderedAccessView * > dxUAVPtrs (2);

	dxConstBufPtrs[0] = _pGroupSizeDxConstBuf;
	d3dUtils.setConstBufs (QD3DUtils::COMPUTE_SHADER, 0, dxConstBufPtrs);

	// Up-Sweep
	d3dUtils.setComputeShader (_pScanUpSweepCS);

	int i = 0;
	for (int totalElemNum = elemNum; totalElemNum > 1;
			totalElemNum = ceil (totalElemNum / (double)GROUP_SIZE), ++i) {
		constVars.var.totalElemNum = totalElemNum;
		totalElemNumArr[i] = totalElemNum;
		d3dUtils.updateSubresource (_pGroupSizeDxConstBuf, 0, &constVars);
		dxUAVPtrs[0] = _tempPackPtrs[i]->getUAV ();
		if (i == _tempPackPtrs.size () - 1)
			dxUAVPtrs[1] = NULL;
		else
			dxUAVPtrs[1] = _tempPackPtrs[i + 1]->getUAV ();
		d3dUtils.setUAVs (QD3DUtils::COMPUTE_SHADER, 0, dxUAVPtrs);
		d3dUtils.dispatch (ceil (totalElemNum / (double)GROUP_SIZE), 1, 1);
		d3dUtils.clearUAVs (QD3DUtils::COMPUTE_SHADER, 0, dxUAVPtrs.size ());
	}

	// Down-Sweep
	d3dUtils.setComputeShader (_pScanDownSweepCS);

	for (i = i - 2; i >= 0; --i) {
		constVars.var.totalElemNum = totalElemNumArr[i];
		d3dUtils.updateSubresource (_pGroupSizeDxConstBuf, 0, &constVars);
		dxSRVPtrs[0] = _tempPackPtrs[i + 1]->getSRV ();
		dxUAVPtrs[0] = _tempPackPtrs[i]->getUAV ();
		dxUAVPtrs[1] = NULL;
		d3dUtils.setSRVs (QD3DUtils::COMPUTE_SHADER, 0, dxSRVPtrs);
		d3dUtils.setUAVs (QD3DUtils::COMPUTE_SHADER, 0, dxUAVPtrs);
		d3dUtils.dispatch (totalElemNumArr[i + 1], 1, 1);
		d3dUtils.clearSRVs (QD3DUtils::COMPUTE_SHADER, 0, dxSRVPtrs.size ());
		d3dUtils.clearUAVs (QD3DUtils::COMPUTE_SHADER, 0, dxUAVPtrs.size ());
	}

	delete [] totalElemNumArr;
}

